import backtrader as bt
import numpy as np
import pandas as pd
from typing import Dict, Any, Optional

class MomentumStrategy(bt.Strategy):
    """
    动量策略
    
    策略逻辑：
    1. 计算价格动量（短期均线vs长期均线）
    2. 成交量确认趋势强度
    3. 在动量向上且成交量放大时买入
    4. 在动量转弱时卖出
    5. 设置移动止损
    """
    
    params = (
        ('fast_period', 5),     # 快速均线周期
        ('slow_period', 20),    # 慢速均线周期
        ('volume_period', 10),  # 成交量均线周期
        ('momentum_threshold', 0.02), # 动量阈值
        ('volume_threshold', 1.5),    # 成交量放大倍数
        ('trailing_stop', 0.08),      # 移动止损比例
        ('position_size', 0.15),      # 仓位比例
    )
    
    def __init__(self):
        # 技术指标
        self.fast_ma = bt.indicators.SMA(
            self.data.close, 
            period=self.params.fast_period
        )
        self.slow_ma = bt.indicators.SMA(
            self.data.close, 
            period=self.params.slow_period
        )
        
        # 动量指标
        self.momentum = bt.indicators.MomentumOscillator(
            self.data.close,
            period=10
        )
        
        # 成交量指标
        self.volume_ma = bt.indicators.SMA(
            self.data.volume,
            period=self.params.volume_period
        )
        
        # 交易状态
        self.order = None
        self.trailing_stop_order = None
        self.entry_price = None
        
    def next(self):
        if self.order:
            return
            
        if not self.position:
            # 寻找买入机会
            momentum_signal = (self.fast_ma[0] > self.slow_ma[0] and 
                             self.fast_ma[-1] <= self.slow_ma[-1])  # 金叉
            
            volume_signal = (self.data.volume[0] > 
                           self.volume_ma[0] * self.params.volume_threshold)
            
            price_momentum = (self.data.close[0] - self.data.close[-5]) / self.data.close[-5]
            
            if (momentum_signal and volume_signal and 
                price_momentum > self.params.momentum_threshold):
                
                # 计算买入数量
                cash = self.broker.getcash()
                price = self.data.close[0]
                size = int((cash * self.params.position_size) / price)
                
                if size > 0:
                    self.order = self.buy(size=size)
                    self.entry_price = price
                    print(f'买入信号: 价格={price:.2f}, 动量={price_momentum:.3f}')
        
        else:
            # 寻找卖出机会
            momentum_weak = (self.fast_ma[0] < self.slow_ma[0] and 
                           self.fast_ma[-1] >= self.slow_ma[-1])  # 死叉
            
            price_momentum = (self.data.close[0] - self.data.close[-5]) / self.data.close[-5]
            momentum_negative = price_momentum < -self.params.momentum_threshold
            
            if momentum_weak or momentum_negative:
                self.order = self.sell(size=self.position.size)
                print(f'卖出信号: 价格={self.data.close[0]:.2f}, 动量={price_momentum:.3f}')
    
    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return
            
        if order.status in [order.Completed]:
            if order.isbuy():
                # 设置移动止损
                self.trailing_stop_order = self.sell(
                    size=order.executed.size,
                    exectype=bt.Order.StopTrail,
                    trailpercent=self.params.trailing_stop * 100
                )
                
                print(f'买入完成: 价格={order.executed.price:.2f}, 数量={order.executed.size}')
            else:
                # 取消移动止损
                if self.trailing_stop_order:
                    self.cancel(self.trailing_stop_order)
                    
                print(f'卖出完成: 价格={order.executed.price:.2f}, 数量={order.executed.size}')
        
        self.order = None

def get_strategy_params() -> Dict[str, Any]:
    """获取策略参数"""
    return {
        'name': '动量策略',
        'description': '基于价格动量和成交量确认的趋势跟踪策略',
        'parameters': {
            'fast_period': {'type': 'int', 'default': 5, 'min': 3, 'max': 10, 'description': '快速均线周期'},
            'slow_period': {'type': 'int', 'default': 20, 'min': 15, 'max': 30, 'description': '慢速均线周期'},
            'volume_period': {'type': 'int', 'default': 10, 'min': 5, 'max': 20, 'description': '成交量均线周期'},
            'momentum_threshold': {'type': 'float', 'default': 0.02, 'min': 0.01, 'max': 0.05, 'description': '动量阈值'},
            'volume_threshold': {'type': 'float', 'default': 1.5, 'min': 1.2, 'max': 2.5, 'description': '成交量放大倍数'},
            'trailing_stop': {'type': 'float', 'default': 0.08, 'min': 0.05, 'max': 0.15, 'description': '移动止损比例'},
            'position_size': {'type': 'float', 'default': 0.15, 'min': 0.1, 'max': 0.3, 'description': '仓位比例'}
        }
    } 